import os
import cv2
import numpy as np

dirpath = 'temp'

def add_alpha_channel(img):
    """
    为jpg图像添加alpha通道, 目的: 统一为png格式
    """

    b_channel, g_channel, r_channel = cv2.split(img)  # 剥离jpg图像通道
    alpha_channel = np.ones(b_channel.shape, dtype=b_channel.dtype) * 255  # 创建Alpha通道

    img_new = cv2.merge((b_channel, g_channel, r_channel, alpha_channel))  # 融合通道
    return img_new


def show_img(img, title="gray"):
    cv2.imshow(title, img)
    # 关闭窗口
    cv2.waitKey()
    cv2.destroyAllWindows()


# 边缘检测
def getCanny(image):
    # 高斯模糊
    binary = cv2.GaussianBlur(image, (3, 3), 2, 2)
    # 边缘检测
    binary = cv2.Canny(binary, 60, 240, apertureSize=3)
    # 膨胀操作，尽量使边缘闭合
    kernel = np.ones((3, 3), np.uint8)
    binary = cv2.dilate(binary, kernel, iterations=1)
    return binary


# 求出面积最大的轮廓
def findMaxContour(image):
    # 寻找边缘
    # _, contours = cv2.findContours(image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)     # contours, hierarchy
    contours, hierarchy = cv2.findContours(image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)     # contours, hierarchy

    # 计算面积
    max_area = 0.0
    max_contour = []
    for contour in contours:
        currentArea = cv2.contourArea(contour)
        if currentArea > max_area:
            max_area = currentArea
            max_contour = contour
    return max_contour, max_area


# 多边形拟合凸包的四个顶点
def getBoxPoint(contour):
    # 多边形拟合凸包
    hull = cv2.convexHull(contour)
    epsilon = 0.02 * cv2.arcLength(contour, True)
    approx = cv2.approxPolyDP(hull, epsilon, True)
    approx = approx.reshape((len(approx), 2))
    return approx


# 适配原四边形点集
def adaPoint(box, pro):
    box_pro = box
    if pro != 1.0:
        box_pro = box/pro
    box_pro = np.trunc(box_pro)
    return box_pro


# 四边形顶点排序，[top-left, top-right, bottom-right, bottom-left]
def orderPoints(pts):
    rect = np.zeros((4, 2), dtype="float32")
    s = pts.sum(axis=1)
    rect[0] = pts[np.argmin(s)]
    rect[2] = pts[np.argmax(s)]
    diff = np.diff(pts, axis=1)
    rect[1] = pts[np.argmin(diff)]
    rect[3] = pts[np.argmax(diff)]
    return rect


# 计算长宽
def pointDistance(a, b):
    return int(np.sqrt(np.sum(np.square(a - b))))


# 透视变换
def warpImage(image, box):
    w, h = pointDistance(box[0], box[1]), \
           pointDistance(box[1], box[2])
    dst_rect = np.array([[0, 0],
                         [w - 1, 0],
                         [w - 1, h - 1],
                         [0, h - 1]], dtype='float32')
    M = cv2.getPerspectiveTransform(box, dst_rect)
    warped = cv2.warpPerspective(image, M, (w, h))
    return warped


# 固定尺寸
def resizeImg(image, height=900):
    h, w = image.shape[:2]
    pro = height / h
    size = (int(w * pro), int(height))
    img = cv2.resize(image, size)
    return img


if __name__ == '__main__':
    imgs = [os.path.join(dirpath, img) for img in os.listdir(dirpath)]
    img_a = imgs[1]

    base_weight, base_height = base_shape = [600, 300]

    img: str = img_a
    img_0 = cv2.imread(img)

    # img = cv2.resize(img_0, base_shape)
    # img = cv2.pyrMeanShiftFiltering(img, 25, 10)
    #
    # # 转换为灰度图
    # # img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    #
    # # 边缘检测
    # binary_img = getCanny(img)
    #
    # # 最大边缘
    # max_contour, max_area = findMaxContour(binary_img)
    #
    # # 获取矩形顶点
    # boxes = getBoxPoint(max_contour)
    #
    # img_0.shape     # hw
    # base_shape      # wh

    # cv2.drawContours(img, max_contour, -1, (255, 0, 0), 3)
    # show_img(img, title='haha')

    # ---------------
    image = cv2.imread(img_a)

    # ratio = 900 / image.shape[0]
    # img = resizeImg(image)

    ratio = base_shape[1] / image.shape[0]
    img = cv2.resize(img_0, base_shape)

    binary_img = getCanny(img)
    # show_img(binary_img, title='haha')

    max_contour, max_area = findMaxContour(binary_img)

    # contours, hierarchy = cv2.findContours(binary_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)  # contours, hierarchy
    cv2.drawContours(img, max_contour, -1, (255, 0, 0), 3)
    show_img(img, title='drawContours')

    boxes = getBoxPoint(max_contour)
    boxes = adaPoint(boxes, ratio)
    boxes = orderPoints(boxes)
    # 透视变化
    warped = warpImage(image, boxes)
    show_img(warped, title='warped')


    # img: np.ndarray
    # img_a.shape

    # image = img_a
    # image = cv2.pyrMeanShiftFiltering(image, 25, 10)
    if 0:
        binary_img
        img.shape
        img_black = np.zeros(img.shape, np.uint8)
        # 浅灰色背景 200
        # img_black.fill(200)
        cv2.drawContours(img_black, max_contour, -1, (255, 0, 0), 3)
        show_img(img_black, 'img_black')

        cv2.drawContours(binary_img, max_contour, -1, (255, 0, 0), 3)
        show_img(binary_img, 'binary_img')

        cnts = cv2.findContours(binary_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cnts = cnts[0] if len(cnts) == 2 else cnts[1]
        for c in cnts:
            peri = cv2.arcLength(c, True)
            approx = cv2.approxPolyDP(c, 0.015 * peri, True)
            if len(approx) == 4:
                x, y, w, h = cv2.boundingRect(approx)
                # (x, y), (w, h), theta = cv2.minAreaRect(approx)
                print(approx)
                cv2.rectangle(img, (x, y), (x + w, y + h), (36, 255, 12), 2)

        show_img(img)

        # 形状包围盒
        bounding_box = cv2.boundingRect(img)
        [x, y, w, h] = bounding_box
        img = cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)
        show_img(img, 'ending')
    1